
//
//Copyright (C) Sasha Goldshtein, 2009
//

#include "MultiProcessorExtensions.h"
#include <msclr\marshal.h>
#include <msclr\marshal_windows.h>

namespace MultiProcessorExtensions
{
	ProcessorGroup::ProcessorGroup(System::UInt16 group, System::Byte processors)
	{
		_groupNumber = group;
		_processorCount = processors;
	}

	System::UInt16 ProcessorGroup::Group::get()
	{
		return _groupNumber;
	}

	System::Byte ProcessorGroup::Processors::get()
	{
		return _processorCount;
	}


	ProcessorNumber::ProcessorNumber(System::UInt16 group, System::Byte processor)
	{
		_groupNumber = group;
		_processorNumber = processor;
	}

	System::UInt16 ProcessorNumber::Group::get()
	{
		return _groupNumber;
	}

	System::Byte ProcessorNumber::Processor::get()
	{
		return _processorNumber;
	}


	System::UInt32 Processor::ActiveProcessorTotalCount::get()
	{
		return ::GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
	}

	System::UInt16 Processor::ActiveProcessorGroupCount::get()
	{
		return ::GetActiveProcessorGroupCount();
	}
	System::UInt32 Processor::GetActiveProcessorCount(System::UInt16 groupNumber)
	{
		return ::GetActiveProcessorCount(groupNumber);
	}

	System::Collections::Generic::IEnumerable<ProcessorGroup>^ Processor::ProcessorGroups::get()
	{
		int groupCount = ActiveProcessorGroupCount;
		System::Collections::Generic::List<ProcessorGroup>^ groups =
			gcnew System::Collections::Generic::List<ProcessorGroup>(groupCount);

		for (int group = 0; group < groupCount; ++group)
		{
			groups->Add(ProcessorGroup(group, GetActiveProcessorCount(group)));
		}
		
		return groups;
	}

	ProcessorNumber Processor::CurrentProcessor::get()
	{
		PROCESSOR_NUMBER procNum = {0};
		::GetCurrentProcessorNumberEx(&procNum);
		return ProcessorNumber(procNum.Group, procNum.Number);
	}

	array<ProcessorRelationship>^ Processor::GetProcessorRelationships()
	{
		DWORD dwLength = 0;
		if (::GetLogicalProcessorInformationEx(::RelationProcessorCore, NULL, &dwLength))
		{
			return gcnew array<ProcessorRelationship>(0);
		}
		DWORD dwLastError = GetLastError();
		if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
		{
			throw gcnew System::ComponentModel::Win32Exception(dwLastError);
		}

		System::Collections::Generic::List<ProcessorRelationship>^ result = 
			gcnew System::Collections::Generic::List<ProcessorRelationship>();
		PBYTE buffer = new BYTE[dwLength];
		try
		{
			if (!::GetLogicalProcessorInformationEx(::RelationProcessorCore, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)buffer, &dwLength))
			{
				throw gcnew System::ComponentModel::Win32Exception(GetLastError());
			}
			for (PBYTE temp = buffer; temp < buffer + dwLength; temp += dwLength)
			{
				PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)temp;
				
				System::Collections::Generic::List<GroupAffinity>^ groups =
					gcnew System::Collections::Generic::List<GroupAffinity>();
				for (int group = 0; group < info->Processor.GroupCount; ++group)
				{
					groups->Add(
						GroupAffinity(
							(System::UIntPtr)(System::UInt64)info->Processor.GroupMask[group].Mask,
							info->Processor.GroupMask[group].Group));
				}
				
				ProcessorRelationship rel(
					info->Processor.Flags == LTP_PC_SMT,
					groups);
				result->Add(rel);
			}
			
		}
		finally
		{
			delete[] buffer;
		}
		return result->ToArray();
	}
	
	array<ProcessorGroupRelationship>^ Processor::GetProcessorGroupRelationships()
	{
		DWORD dwLength = 0;
		if (::GetLogicalProcessorInformationEx(::RelationGroup, NULL, &dwLength))
		{
			return gcnew array<ProcessorGroupRelationship>(0);
		}
		DWORD dwLastError = GetLastError();
		if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
		{
			throw gcnew System::ComponentModel::Win32Exception(dwLastError);
		}

		System::Collections::Generic::List<ProcessorGroupRelationship>^ result = 
			gcnew System::Collections::Generic::List<ProcessorGroupRelationship>();
		PBYTE buffer = new BYTE[dwLength];
		try
		{
			if (!::GetLogicalProcessorInformationEx(::RelationGroup, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)buffer, &dwLength))
			{
				throw gcnew System::ComponentModel::Win32Exception(GetLastError());
			}
			for (PBYTE temp = buffer; temp < buffer + dwLength; temp += dwLength)
			{
				PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)temp;
				
				System::Collections::Generic::List<ProcessorGroupInfo>^ groupInfos =
					gcnew System::Collections::Generic::List<ProcessorGroupInfo>();
				for (int group = 0; group < info->Group.ActiveGroupCount; ++group)
				{
					groupInfos->Add(
						ProcessorGroupInfo(
							info->Group.GroupInfo[group].MaximumProcessorCount,
							info->Group.GroupInfo[group].ActiveProcessorCount,
							(System::UIntPtr)(System::UInt64)info->Group.GroupInfo[group].ActiveProcessorMask));
				}
				
				ProcessorGroupRelationship rel(
					info->Group.MaximumGroupCount,
					info->Group.ActiveGroupCount,
					groupInfos);
				result->Add(rel);
			}
			
		}
		finally
		{
			delete[] buffer;
		}
		return result->ToArray();
	}

	array<NumaNodeRelationship>^ Processor::GetNumaNodeRelationships()
	{
		DWORD dwLength = 0;
		if (::GetLogicalProcessorInformationEx(::RelationNumaNode, NULL, &dwLength))
		{
			return gcnew array<NumaNodeRelationship>(0);
		}
		DWORD dwLastError = GetLastError();
		if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
		{
			throw gcnew System::ComponentModel::Win32Exception(dwLastError);
		}

		System::Collections::Generic::List<NumaNodeRelationship>^ result = 
			gcnew System::Collections::Generic::List<NumaNodeRelationship>();
		PBYTE buffer = new BYTE[dwLength];
		try
		{
			if (!::GetLogicalProcessorInformationEx(::RelationNumaNode, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)buffer, &dwLength))
			{
				throw gcnew System::ComponentModel::Win32Exception(GetLastError());
			}
			for (PBYTE temp = buffer; temp < buffer + dwLength; temp += dwLength)
			{
				PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)temp;
				NumaNodeRelationship rel(
					info->NumaNode.NodeNumber,
					GroupAffinity((System::UIntPtr)(System::UInt64)info->NumaNode.GroupMask.Mask, info->NumaNode.GroupMask.Group));
				result->Add(rel);
			}
			
		}
		finally
		{
			delete[] buffer;
		}
		return result->ToArray();
	}

	array<CacheRelationship>^ Processor::GetCacheRelationships()
	{
		DWORD dwLength = 0;
		if (::GetLogicalProcessorInformationEx(::RelationCache, NULL, &dwLength))
		{
			return gcnew array<CacheRelationship>(0);
		}
		DWORD dwLastError = GetLastError();
		if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
		{
			throw gcnew System::ComponentModel::Win32Exception(dwLastError);
		}

		System::Collections::Generic::List<CacheRelationship>^ result = 
			gcnew System::Collections::Generic::List<CacheRelationship>();
		PBYTE buffer = new BYTE[dwLength];
		try
		{
			if (!::GetLogicalProcessorInformationEx(::RelationCache, (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)buffer, &dwLength))
			{
				throw gcnew System::ComponentModel::Win32Exception(GetLastError());
			}
			for (PBYTE temp = buffer; temp < buffer + dwLength; temp += dwLength)
			{
				PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX info = (PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX)temp;
				CacheRelationship rel(
					info->Cache.Level,
					info->Cache.Associativity,
					info->Cache.LineSize,
					info->Cache.CacheSize,
					(ProcessorCacheType)info->Cache.Type,
					GroupAffinity((System::UIntPtr)(System::UInt64)info->Cache.GroupMask.Mask, info->Cache.GroupMask.Group));
				result->Add(rel);
			}
			
		}
		finally
		{
			delete[] buffer;
		}
		return result->ToArray();		
	}

	System::UInt16 Processor::MaximumProcessorGroupCount::get()
	{
		return ::GetMaximumProcessorGroupCount();
	}

	System::UInt32 Processor::GetMaximumProcessorCount(System::UInt16 groupNumber)
	{
		return ::GetMaximumProcessorCount(groupNumber);
	}


	GroupAffinity::GroupAffinity(System::UIntPtr mask, System::UInt16 group)
	{
		Mask = mask;
		Group = group;
		Reserved = 0;
	}
	void GroupAffinity::Set(System::Byte processorNumber)
	{
		System::UInt64 mask = Mask.ToUInt64();
		mask |= (1LL << processorNumber);
		Mask = (System::UIntPtr)mask;
	}
	void GroupAffinity::Clear(System::Byte processorNumber)
	{
		System::UInt64 mask = Mask.ToUInt64();
		mask &= ~(1LL << processorNumber);
		Mask = (System::UIntPtr)mask;
	}
	void GroupAffinity::ClearAll()
	{
		Mask = (System::UIntPtr)(System::UInt32)0;
	}


	ProcessorRelationship::ProcessorRelationship(
		bool coreHasMultipleLogicalProcessors,
		System::Collections::Generic::IEnumerable<GroupAffinity>^ groupMask)
	{
		_coreHasMultipleLogicalProcessors = coreHasMultipleLogicalProcessors;
		_groupMask = gcnew System::Collections::Generic::List<GroupAffinity>(groupMask);
	}

	bool ProcessorRelationship::CoreHasMultipleLogicalProcessors::get()
	{
		return _coreHasMultipleLogicalProcessors;
	}
	System::Collections::Generic::IEnumerable<GroupAffinity>^ ProcessorRelationship::GroupMask::get()
	{
		return _groupMask;
	}


	NumaNodeRelationship::NumaNodeRelationship(
		System::UInt32 nodeNumber,
		GroupAffinity groupAffinity)
	{
		_nodeNumber = nodeNumber;
		_groupAffinity = groupAffinity;
	}

	System::UInt32 NumaNodeRelationship::NodeNumber::get()
	{
		return _nodeNumber;
	}
	GroupAffinity NumaNodeRelationship::Affinity::get()
	{
		return _groupAffinity;
	}


	CacheRelationship::CacheRelationship(
		System::Byte level,
		System::Byte associativity,
		System::UInt16 lineSize,
		System::UInt32 cacheSize,
		ProcessorCacheType type,
		GroupAffinity groupAffinity)
	{
		_level = level;
		_associativity = associativity;
		_lineSize = lineSize;
		_cacheSize = cacheSize;
		_type = type;
		_groupAffinity = groupAffinity;
	}

	System::Byte CacheRelationship::Level::get()
	{
		return _level;
	}

	System::Byte CacheRelationship::Associativity::get()
	{
		return _associativity;
	}
	
	bool CacheRelationship::IsFullyAssociative::get()
	{
		return _associativity == CACHE_FULLY_ASSOCIATIVE;
	}
	
	System::UInt16 CacheRelationship::LineSize::get()
	{
		return _lineSize;
	}

	System::UInt32 CacheRelationship::CacheSize::get()
	{
		return _cacheSize;
	}
	
	ProcessorCacheType CacheRelationship::Type::get()
	{
		return _type;
	}

	GroupAffinity CacheRelationship::Affinity::get()
	{
		return _groupAffinity;
	}


	ProcessorGroupInfo::ProcessorGroupInfo(
		System::Byte maxProcessorCount,
		System::Byte activeProcessorCount,
		System::UIntPtr activeProcessors)
	{
		_maxProcessorCount = maxProcessorCount;
		_activeProcessorCount = activeProcessorCount;
		_activeProcessors = activeProcessors;
	}

	System::Byte ProcessorGroupInfo::MaximumProcessorCount::get()
	{
		return _maxProcessorCount;
	}
	
	System::Byte ProcessorGroupInfo::ActiveProcessorCount::get()
	{
		return _activeProcessorCount;
	}

	System::UIntPtr ProcessorGroupInfo::ActiveProcessors::get()
	{
		return _activeProcessors;
	}

	bool ProcessorGroupInfo::IsProcessorActive(System::Byte index)
	{
		System::UInt64 mask = _activeProcessors.ToUInt64();
		return (mask & (1LL << index)) == 1;
	}


	ProcessorGroupRelationship::ProcessorGroupRelationship(
		System::UInt16 maxGroupCount,
		System::UInt16 activeGroupCount,
		System::Collections::Generic::IEnumerable<ProcessorGroupInfo>^ processors)
	{
		_maxGroupCount = maxGroupCount;
		_activeGroupCount = activeGroupCount;
		_processors = gcnew System::Collections::Generic::List<ProcessorGroupInfo>(processors);
	}

	System::UInt16 ProcessorGroupRelationship::MaximumGroupCount::get()
	{
		return _maxGroupCount;
	}
	
	System::UInt16 ProcessorGroupRelationship::ActiveGroupCount::get()
	{
		return _activeGroupCount;
	}
	
	System::Collections::Generic::IEnumerable<ProcessorGroupInfo>^ ProcessorGroupRelationship::Processors::get()
	{
		return _processors;
	}


	System::UInt64 Numa::GetAvailableMemory(System::UInt16 numaNode)
	{
		ULONGLONG available;
		if (!::GetNumaAvailableMemoryNodeEx(numaNode, &available))
		{
			throw gcnew System::ComponentModel::Win32Exception(GetLastError());
		}
		return available;
	}

	GroupAffinity Numa::GetProcessorMask(System::UInt16 numaNode)
	{
		GROUP_AFFINITY procMask;
		if (!::GetNumaNodeProcessorMaskEx(numaNode, &procMask))
		{
			throw gcnew System::ComponentModel::Win32Exception(GetLastError());
		}
		return GroupAffinity((System::UIntPtr)(System::UInt64)procMask.Mask, procMask.Group);
	}

	System::UInt16 Numa::GetProcessorNode(ProcessorNumber processor)
	{
		USHORT nodeNum;
		PROCESSOR_NUMBER procNum = {0};
		procNum.Group = processor.Group;
		procNum.Number = processor.Processor;
		if (!::GetNumaProcessorNodeEx(&procNum, &nodeNum))
		{
			throw gcnew System::ComponentModel::Win32Exception(GetLastError());
		}
		return nodeNum;
	}
	
	System::UInt16 Numa::GetProximityNode(System::UInt32 proximityId)
	{
		USHORT nodeNum;
		if (!::GetNumaProximityNodeEx(proximityId, &nodeNum))
		{
			throw gcnew System::ComponentModel::Win32Exception(GetLastError());
		}
		return nodeNum;
	}


	void Thread::SetThreadGroupAffinity(
		System::Diagnostics::ProcessThread^ thread,
		GroupAffinity affinity)
	{
		SetThreadGroupAffinity(thread->Id, affinity);
	}

	void Thread::SetThreadGroupAffinity(
		System::UInt32 threadId,
		GroupAffinity affinity)
	{
		HANDLE hThread = ::OpenThread(THREAD_SET_INFORMATION, FALSE, threadId);
		try
		{
			SetThreadGroupAffinity((System::IntPtr)hThread, affinity);
		}
		finally
		{
			::CloseHandle(hThread);
		}
	}
	
	System::UInt32 Thread::CurrentThreadId::get()
	{
		return ::GetCurrentThreadId();
	}

	void Thread::SetThreadGroupAffinity(
		System::IntPtr thread,
		GroupAffinity affinity)
	{
		GROUP_AFFINITY gAffinity = {0};
		gAffinity.Group = affinity.Group;
		gAffinity.Mask = (KAFFINITY)affinity.Mask.ToPointer();

		if (!::SetThreadGroupAffinity((HANDLE)thread.ToPointer(), &gAffinity, NULL))
		{
			throw gcnew System::ComponentModel::Win32Exception(::GetLastError());
		}
	}

	GroupAffinity Thread::GetThreadGroupAffinity(
		System::Diagnostics::ProcessThread^ thread)
	{
		return GetThreadGroupAffinity(thread->Id);
	}
	
	GroupAffinity Thread::GetThreadGroupAffinity(
		System::UInt32 threadId)
	{
		HANDLE hThread = ::OpenThread(THREAD_SET_INFORMATION, FALSE, threadId);
		try
		{
			return GetThreadGroupAffinity((System::IntPtr)hThread);
		}
		finally
		{
			::CloseHandle(hThread);
		}
	}
	
	GroupAffinity Thread::GetThreadGroupAffinity(
		System::IntPtr thread)
	{
		GROUP_AFFINITY gAffinity;
		if (!::GetThreadGroupAffinity((HANDLE)thread.ToPointer(), &gAffinity))
		{
			throw gcnew System::ComponentModel::Win32Exception(::GetLastError());
		}
		return GroupAffinity((System::UIntPtr)(System::UInt64)gAffinity.Mask, gAffinity.Group);
	}


	System::Diagnostics::Process^ Process::CreateWithParentAffinity(
		System::String^ fileName,
		System::String^ arguments)
	{
		msclr::interop::marshal_context context;

		STARTUPINFO startupInfo = {0};
		startupInfo.cb = sizeof(startupInfo);
		PROCESS_INFORMATION processInfo = {0};

		if (!::CreateProcess(
			context.marshal_as<LPCWSTR>(fileName),
			context.marshal_as<LPWSTR>(fileName + " " + arguments),
			NULL,
			NULL,
			FALSE,
			INHERIT_PARENT_AFFINITY,
			NULL,
			NULL,
			&startupInfo,
			&processInfo))
		{
			throw gcnew System::ComponentModel::Win32Exception(GetLastError());
		}

		::CloseHandle(processInfo.hProcess);
		::CloseHandle(processInfo.hThread);

		return System::Diagnostics::Process::GetProcessById(processInfo.dwProcessId);

	}

	array<System::UInt16>^ Process::GetProcessGroupAffinity(
		System::Diagnostics::Process^ process)
	{
		return GetProcessGroupAffinity(process->Id);
	}
	
	array<System::UInt16>^ Process::GetProcessGroupAffinity(
		System::UInt32 processId)
	{
		HANDLE hProcess = ::OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, processId);
		try
		{
			return GetProcessGroupAffinity((System::IntPtr)hProcess);
		}
		finally
		{
			::CloseHandle(hProcess);
		}
	}

	array<System::UInt16>^ Process::GetProcessGroupAffinity(
		System::IntPtr process)
	{
		USHORT uCount = 0;
		if (::GetProcessGroupAffinity((HANDLE)process.ToPointer(), &uCount, NULL))
		{
			return gcnew array<System::UInt16>(0);
		}
		DWORD dwLastError = GetLastError();
		if (dwLastError != ERROR_INSUFFICIENT_BUFFER)
		{
			throw gcnew System::ComponentModel::Win32Exception(dwLastError);
		}

		PUSHORT group = new USHORT[uCount];
		try
		{
			if (!::GetProcessGroupAffinity((HANDLE)process.ToPointer(), &uCount, group))
			{
				throw gcnew System::ComponentModel::Win32Exception(GetLastError());
			}
			array<System::UInt16>^ result = gcnew array<System::UInt16>(uCount);
			for (int i = 0; i < uCount; ++i)
			{
				result[i] = group[i];
			}
			return result;
		}
		finally
		{
			delete[] group;
		}
	}
}